// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions {
namespace {

class MapKeysByTopNValuesTest : public test::FunctionBaseTest {};

TEST_F(MapKeysByTopNValuesTest, emptyMap) {
  RowVectorPtr input = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>({
          "{}",
      }),
  });

  assertEqualVectors(
      evaluate("map_keys_by_top_n_values(c0, 3)", input),
      make_array_vector_from_json<int32_t>({
          "[]",
      }));
}

TEST_F(MapKeysByTopNValuesTest, basic) {
  auto data = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>({
          "{1:3, 2:5, 3:1, 4:4, 5:2}",
          "{1:3, 2:5, 3:null, 4:4, 5:2}",
          "{1:null, 2:null, 3:1, 4:4, 5:null}",
          "{1:10, 2:7, 3:11, 5:4}",
          "{1:10, 2:10, 3:10, 4:10, 5:10}",
          "{1:10, 2:7, 3:0}",
          "{1:null, 2:10}",
          "{}",
          "{1:null, 2:null, 3:null}",
      }),
  });

  auto result = evaluate("map_keys_by_top_n_values(c0, 3)", data);

  auto expected = make_array_vector_from_json<int32_t>({
      "[2, 4, 1]",
      "[2, 4, 1]",
      "[4, 3, 5]",
      "[3, 1, 2]",
      "[5, 4, 3]",
      "[1, 2, 3]",
      "[2, 1]",
      "[]",
      "[3, 2, 1]",
  });

  assertEqualVectors(expected, result);

  // n = 0. Expect empty maps.
  result = evaluate("map_keys_by_top_n_values(c0, 0)", data);

  expected = make_array_vector_from_json<int32_t>({
      "[]",
      "[]",
      "[]",
      "[]",
      "[]",
      "[]",
      "[]",
      "[]",
      "[]",
  });

  assertEqualVectors(expected, result);

  // n is negative. Expect an error.
  POLLUX_ASSERT_THROW(
      evaluate("map_keys_by_top_n_values(c0, -1)", data),
      "n must be greater than or equal to 0");
}

TEST_F(MapKeysByTopNValuesTest, complexKeys) {
  RowVectorPtr input =
      make_row_vector({make_map_vector_from_json<std::string, int64_t>(
          {R"( {"x":1, "y":2} )",
           R"( {"x":1, "x2":-2} )",
           R"( {"ac":1, "cc":3, "dd": 4} )"})});

  assertEqualVectors(
      evaluate("map_keys_by_top_n_values(c0, 1)", input),
      make_array_vector_from_json<std::string>({
          "[\"y\"]",
          "[\"x\"]",
          "[\"dd\"]",
      }));
}

TEST_F(MapKeysByTopNValuesTest, tryFuncWithInputHasNullInArray) {
  auto array_vector = make_nullable_flat_vector<std::int32_t>({1, 2, 3});
  auto flat_vector = make_nullable_array_vector<bool>(
      {{true, std::nullopt}, {true, std::nullopt}, {false, std::nullopt}});
  auto mapvector = make_map_vector(
      /*offsets=*/{0},
      /*keyVector=*/array_vector,
      /*value_vector=*/flat_vector);
  auto rst = evaluate(
      "try(map_keys_by_top_n_values(c0, 6455219767830808341))",
      make_row_vector({mapvector}));
  assert(rst->size() == 1);
  assert(rst->is_null_at(0));

  array_vector = make_nullable_flat_vector<std::int32_t>({1, 2, 3, -1});
  flat_vector = make_nullable_array_vector<std::int32_t>(
      {{1}, {1, 2}, {1, 3, 4}, {4, std::nullopt}});
  mapvector = make_map_vector(
      /*offsets=*/{0},
      /*keyVector=*/array_vector,
      /*value_vector=*/flat_vector);
  auto result = evaluate(
      "try(map_keys_by_top_n_values(c0, 6455219767830808341))",
      make_row_vector({mapvector}));
  auto expected = make_array_vector_from_json<int32_t>({"[-1, 3, 2, 1]"});
  assertEqualVectors(expected, result);
}

TEST_F(MapKeysByTopNValuesTest, complexKeysWithLargeK) {
  RowVectorPtr input =
      make_row_vector({make_map_vector_from_json<std::string, int64_t>(
          {R"( {"x":1, "y":2} )",
           R"( {"x":1, "x2":-2} )",
           R"( {"ac":1, "cc":3, "dd": 4} )"})});

  assertEqualVectors(
      evaluate("map_keys_by_top_n_values(c0, 706100916841560005)", input),
      make_array_vector_from_json<std::string>({
          "[\"y\", \"x\"]",
          "[\"x\", \"x2\"]",
          "[\"dd\", \"cc\", \"ac\"]",
      }));
}

TEST_F(MapKeysByTopNValuesTest, timestampWithTimeZone) {
  auto testMapTopNKeys = [&](const std::vector<int64_t>& keys,
                             const std::vector<int32_t>& values,
                             const std::vector<int64_t>& expectedKeys) {
    const auto map = make_map_vector(
        {0},
        make_flat_vector(keys, TIMESTAMP_WITH_TIME_ZONE()),
        make_flat_vector(values));
    const auto expected = make_array_vector(
        {0}, make_flat_vector(expectedKeys, TIMESTAMP_WITH_TIME_ZONE()));

    const auto result =
        evaluate("map_keys_by_top_n_values(c0, 3)", make_row_vector({map}));

    assertEqualVectors(expected, result);
  };

  testMapTopNKeys(
      {pack(1, 1), pack(2, 2), pack(3, 3), pack(4, 4), pack(5, 5)},
      {3, 5, 1, 4, 2},
      {pack(2, 2), pack(4, 4), pack(1, 1)});
  testMapTopNKeys(
      {pack(5, 1), pack(4, 2), pack(3, 3), pack(2, 4), pack(1, 5)},
      {3, 5, 1, 4, 2},
      {pack(4, 2), pack(2, 4), pack(5, 1)});
  testMapTopNKeys(
      {pack(3, 1), pack(5, 2), pack(1, 3), pack(4, 4), pack(2, 5)},
      {1, 2, 3, 4, 5},
      {pack(2, 5), pack(4, 4), pack(1, 3)});
  testMapTopNKeys(
      {pack(3, 5), pack(5, 4), pack(4, 2), pack(2, 1)},
      {3, 3, 3, 3},
      {pack(5, 4), pack(4, 2), pack(3, 5)});
  testMapTopNKeys({}, {}, {});
}
} // namespace
} // namespace kumo::pollux::functions
